import yaml

from BaseParams import BaseParams
from constants import AlgoType
from data_process.DatasetGenerator import DatasetGenerator
from implement.model_distillation import LipReading


class MultiDopplerParams(BaseParams):
    def __init__(self):
        super(MultiDopplerParams, self).__init__()
        self.sonicStart = 17000
        self.sonicInter = 500
        self.sonicNum = 8
        self.volunList = [1]
        self.innerDirName = "mix_17k_500"
        self.algoType = AlgoType.MULTI_DOPPLER

datasetGenerator = DatasetGenerator(MultiDopplerParams())
datasetGenerator.emsembleProcess()

params = MultiDopplerParams()

yamlPath = 'config/conv.yaml'
file = open(yamlPath, 'r')
paras = yaml.load(file, Loader=yaml.FullLoader)
file.close()

# COMTYPE = ['alexnet','vgg']
COMTYPE = ['resnet8']
ABLATION = ['resnet', 'kl', 'attn', 'kl_attn']
modelAllType = ['resnet8']
modelAllFun = ['base']
for modelType in modelAllType:
    for modelFunc in modelAllFun:
        paras['modelType'] = modelType
        paras['modelFunc'] = modelFunc
        paras['curVolun'] = '-1'
        paras['trainRate'] = 1
        lipModel = LipReading(paras)
        lipModel.Train()